import codecs

prompt_path = '/home/rickwwang/project_research/fairseq/data-bin/prompt.txt.test'
all_prompts = codecs.open(prompt_path, 'r', encoding='utf8').readlines()
all_prompts = map(lambda x: x.strip(), all_prompts)

exp_file_src = []
exp_file_src.append('fire_data/running/abs2story_small_ls0/41src.txt')
exp_file_src.append('fire_data/running/abs2story_small_ls0/41src.txt.own.prompt2abs')

exp_file_tgt = []
exp_file_tgt.append('fire_data/running/abs2story_small_ls0/41hyp.txt.1.topk20.t1.0')
exp_file_tgt.append('fire_data/running/abs2story_small_ls0/41hyp.txt.own.prompt2abs.1.topk20.t1.0')

base_path = '/home/rickwwang/project_research/'

for exp1, exp2 in zip(exp_file_src, exp_file_tgt):
    story_prompt_path = base_path + exp1
    story_path = base_path + exp2
    result_path = story_path + '.abs.human.test'

    all_stories_prompts = codecs.open(story_prompt_path, 'r', encoding='utf8').readlines()
    all_stories = codecs.open(story_path, 'r', encoding='utf8').readlines()
    all_stories = zip(all_stories_prompts, all_stories)

    promp_to_story = dict()
    for s in all_stories:
        item = s[0].strip()
        # if '<SEP>' in key:
        key = item.split(' <SEP> ')[0]
        value = item.split(' <SEP> ')[1]
        if key in promp_to_story:
            promp_to_story[key].append(value)
        else:
            promp_to_story[key] = [value]

    format_stories = []
    for p in all_prompts:
        if p in promp_to_story:
            ss = promp_to_story[p]

            format_stories.append(p)
            format_stories.append(ss[0])
        else:
            format_stories.append('null')
            format_stories.append('null')

    with codecs.open(result_path, 'w', encoding='utf8') as fout:
        fout.write('\n'.join(format_stories))
        fout.write('\n')